import os
import time
from collections import deque
import pickle
import cv2
import json
import rclpy
from rclpy.time import Time
from rclpy.duration import Duration
from rclpy.executors import ExternalShutdownException
from rclpy.executors import MultiThreadedExecutor
import tf2_ros
from cv_bridge import CvBridge
import sensor_msgs_py.point_cloud2 as pc2
from semantic_nav.slam_utils import (points_to_depth_image,
                                     render_color_points,
                                     transform_to_Rt,
                                     create_point_cloud_from_rgb_depth,
                                     numpy_to_o3d_pcd,
                                     preprocess_point_cloud)
from semantic_nav.ros_utils import ColorDepthSyncer
from semantic_nav.topomap_locator import TopoMapLocator
from semantic_nav.log_utils import get_logger
logger = get_logger()
package_name = 'semantic_nav'


class TopoMapBuilder(ColorDepthSyncer):
    def __init__(self, node_name='topmap_builder', namespace='', voxel_size=0.1):
        super().__init__(node_name=node_name, namespace=namespace)
        self.declare_parameter("keyframe_period", 0.5)
        self.declare_parameter("node_dist_thre", 1.0)
        self.declare_parameter("base_frame", "base_link")
        self.declare_parameter("map_frame", "map")
        
        self.tf_buffer = tf2_ros.Buffer(cache_time=Duration(seconds=20.0))
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self, spin_thread=True)
        self.cv_bridge = CvBridge()
        self.topomap = None
        self.last_time = None
        self.time_stamp = None
        self.voxel_size=voxel_size

    def init(self, path=None, queue_size=10, slop=None):
        self.keyframe_period = self.get_parameter("keyframe_period").value
        self.node_dist_thre = self.get_parameter("node_dist_thre").value
        self.base_frame = self.get_parameter("base_frame").value
        self.map_frame = self.get_parameter("map_frame").value
        self.path = path
        self.topomap = TopoMapLocator(path)
        self.info_path = os.path.join(self.path, 'camera.pkl')
        if os.path.exists(self.info_path):
            with open(self.info_path, 'rb') as f:
                save_dict = pickle.load(f)
                self.camera_intrinsics = save_dict['camera_intrinsics']
        else:
            self.camera_intrinsics = None
        super().init(queue_size=queue_size, slop=slop)
        return self

    def destroy_node(self):
        with open(self.info_path, 'wb') as f:
            save_dict = {"num_color": self.num_color,
                         "camera_intrinsics": self.camera_intrinsics}
            pickle.dump(save_dict, f, protocol=pickle.HIGHEST_PROTOCOL)
        del self.topomap
        super().destroy_node()

    def process_observation(self, color_msgs, depth_msgs):
        self.get_logger().info(f"processing observation")
        start_time = time.time()
        time_stamp, self.observations = self.convert_observation(color_msgs, depth_msgs)
        end_time = time.time()
        logger.info(f"convert observation consumes time: {end_time - start_time:.3f}s")
        this_time = Time().from_msg(time_stamp)

        try:
            tf = self.tf_buffer.lookup_transform(self.map_frame, self.base_frame, this_time, Duration(seconds=1.0))
            this_pose = transform_to_Rt(tf)
            self.topomap.locate_node(self.observations, this_pose)
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            logger.error(f"Error looking up transform: {e}")
        self.last_time = this_time

    def convert_observation(self, color_msgs, depth_msgs):
        color_images = []
        depth_images = []
        points_array = []
        for i in range(len(color_msgs)):
            cv_color_image = self.cv_bridge.imgmsg_to_cv2(color_msgs[i], desired_encoding="bgr8")
            color_images.append(cv_color_image)

        for i in range(len(depth_msgs)):

            if self.use_point_cloud:
                try:
                    tf = self.tf_buffer.lookup_transform(color_msgs[i].header.frame_id, depth_msgs[i].header.frame_id, Time(), Duration(seconds=1.0))
                    R, t = transform_to_Rt(tf)
                except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
                    logger.error(f"Error looking up transform: {e}")
                points = pc2.read_points_numpy(depth_msgs[i], field_names=["x", "y", "z"], skip_nans=True)
                points = (R @ points.T).T + t.reshape(1, 3)
                cv_depth_image = points_to_depth_image(color_images[i], points, self.camera_intrinsics[i])
                colored_points = render_color_points(color_images[i], points, self.camera_intrinsics[i])
                pcd = numpy_to_o3d_pcd(colored_points, xyz_only=False)
            else:
                cv_depth_image = self.cv_bridge.imgmsg_to_cv2(depth_msgs[i], desired_encoding="32FC1")
                pcd = create_point_cloud_from_rgb_depth(cv_color_image, cv_depth_image, self.camera_intrinsics[i])

            pcd = preprocess_point_cloud(pcd, self.voxel_size)
            points_array.append(pcd)
            depth_images.append(cv_depth_image)

        return color_msgs[0].header.stamp, list(zip(color_images, depth_images, points_array))


def main(args=None):
    rclpy.init(args=args)
    try:
        from ament_index_python.packages import get_package_prefix
        map_dir = os.path.join(get_package_prefix(package_name), 'maps')
        topomap_builder = TopoMapBuilder()
        
        topomap_builder.set_parameters([rclpy.parameter.Parameter(name="color_topics", value=["/zed/zed_node/rgb/image_rect_color"]),
                                        rclpy.parameter.Parameter(name="depth_topics", value=["/zed/zed_node/depth/depth_registered"]),
                                        rclpy.parameter.Parameter(name="color_info_topics", value=["/zed/zed_node/rgb/camera_info"]),
                                        rclpy.parameter.Parameter(name="depth_info_topics", value=["/zed/zed_node/depth/camera_info"])])

        topomap_builder.init(map_dir)
        executor = MultiThreadedExecutor()
        rclpy.spin(topomap_builder, executor=executor)
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        pass
    finally:
        topomap_builder.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
